import os
import pandas as pd
import json
import logging
from datetime import datetime

logger = logging.getLogger(__name__)

class DataCacheManager:
    """
    数据缓存管理器，用于管理股票数据的缓存存储和读取
    """
    def __init__(self, cache_dir=None):
        """
        初始化缓存管理器
        
        Parameters:
        -----------
        cache_dir : str, optional
            缓存目录路径，默认为项目根目录下的data_cache文件夹
        """
        if cache_dir is None:
            self.cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data_cache')
        else:
            self.cache_dir = cache_dir
            
        # 确保缓存目录存在
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)
            logger.info(f"创建数据缓存目录: {self.cache_dir}")
            
        # 创建子目录
        self.stock_data_dir = os.path.join(self.cache_dir, 'stock_data')
        self.trade_dates_dir = os.path.join(self.cache_dir, 'trade_dates')
        self.limit_data_dir = os.path.join(self.cache_dir, 'limit_data')
        self.index_data_dir = os.path.join(self.cache_dir, 'index_data')
        
        # 数据预加载缓存
        self.preloaded_stock_data = {}
        self.preloaded_future_returns = {}
        self.preloaded_index_data = {}
        self.preloaded_limit_data = {}
        
        for directory in [self.stock_data_dir, self.trade_dates_dir, self.limit_data_dir, self.index_data_dir]:
            if not os.path.exists(directory):
                os.makedirs(directory)
    
    def get_stock_data_cache_path(self, trade_date):
        """
        获取指定交易日股票数据的缓存文件路径
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
            
        Returns:
        --------
        str
            缓存文件路径
        """
        return os.path.join(self.stock_data_dir, f"stock_data_{trade_date}.csv")
    
    def get_trade_dates_cache_path(self, start_date, end_date):
        """
        获取指定日期范围内交易日历的缓存文件路径
        
        Parameters:
        -----------
        start_date : str
            开始日期，格式为YYYYMMDD
        end_date : str
            结束日期，格式为YYYYMMDD
            
        Returns:
        --------
        str
            缓存文件路径
        """
        return os.path.join(self.trade_dates_dir, f"trade_dates_{start_date}_{end_date}.json")
    
    def get_limit_data_cache_path(self, trade_date):
        """
        获取指定交易日涨停数据的缓存文件路径
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
            
        Returns:
        --------
        str
            缓存文件路径
        """
        return os.path.join(self.limit_data_dir, f"limit_data_{trade_date}.csv")
    
    def get_index_data_cache_path(self, trade_date):
        """
        获取指定交易日指数数据的缓存文件路径
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
            
        Returns:
        --------
        str
            缓存文件路径
        """
        return os.path.join(self.index_data_dir, f"index_data_{trade_date}.csv")
    
    def cache_exists(self, cache_path):
        """
        检查缓存文件是否存在
        
        Parameters:
        -----------
        cache_path : str
            缓存文件路径
            
        Returns:
        --------
        bool
            缓存文件是否存在
        """
        return os.path.exists(cache_path)
    
    def save_stock_data(self, trade_date, data):
        """
        保存股票数据到缓存
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
        data : pandas.DataFrame
            股票数据
            
        Returns:
        --------
        bool
            保存是否成功
        """
        try:
            cache_path = self.get_stock_data_cache_path(trade_date)
            data.to_csv(cache_path, index=False)
            logger.debug(f"股票数据已缓存: {cache_path}")
            return True
        except Exception as e:
            logger.error(f"保存股票数据缓存失败: {str(e)}")
            return False
    
    def load_stock_data(self, trade_date):
        """
        从缓存加载股票数据
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
            
        Returns:
        --------
        pandas.DataFrame or None
            股票数据，如果缓存不存在则返回None
        """
        cache_path = self.get_stock_data_cache_path(trade_date)
        if self.cache_exists(cache_path):
            try:
                data = pd.read_csv(cache_path)
                # 转换日期列为datetime类型
                if 'list_date' in data.columns:
                    data['list_date'] = pd.to_datetime(data['list_date'])
                logger.debug(f"从缓存加载股票数据: {cache_path}")
                return data
            except Exception as e:
                logger.error(f"加载股票数据缓存失败: {str(e)}")
                return None
        return None
    
    def save_trade_dates(self, start_date, end_date, trade_dates):
        """
        保存交易日历到缓存
        
        Parameters:
        -----------
        start_date : str
            开始日期，格式为YYYYMMDD
        end_date : str
            结束日期，格式为YYYYMMDD
        trade_dates : list
            交易日期列表
            
        Returns:
        --------
        bool
            保存是否成功
        """
        try:
            cache_path = self.get_trade_dates_cache_path(start_date, end_date)
            with open(cache_path, 'w') as f:
                json.dump(trade_dates, f)
            logger.debug(f"交易日历已缓存: {cache_path}")
            return True
        except Exception as e:
            logger.error(f"保存交易日历缓存失败: {str(e)}")
            return False
    
    def load_trade_dates(self, start_date, end_date):
        """
        从缓存加载交易日历
        
        Parameters:
        -----------
        start_date : str
            开始日期，格式为YYYYMMDD
        end_date : str
            结束日期，格式为YYYYMMDD
            
        Returns:
        --------
        list or None
            交易日期列表，如果缓存不存在则返回None
        """
        cache_path = self.get_trade_dates_cache_path(start_date, end_date)
        if self.cache_exists(cache_path):
            try:
                with open(cache_path, 'r') as f:
                    trade_dates = json.load(f)
                logger.debug(f"从缓存加载交易日历: {cache_path}")
                return trade_dates
            except Exception as e:
                logger.error(f"加载交易日历缓存失败: {str(e)}")
                return None
        return None
    
    def save_limit_data(self, trade_date, data):
        """
        保存涨停数据到缓存
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
        data : pandas.DataFrame
            涨停数据
            
        Returns:
        --------
        bool
            保存是否成功
        """
        try:
            cache_path = self.get_limit_data_cache_path(trade_date)
            data.to_csv(cache_path, index=False)
            logger.debug(f"涨停数据已缓存: {cache_path}")
            return True
        except Exception as e:
            logger.error(f"保存涨停数据缓存失败: {str(e)}")
            return False
    
    def load_limit_data(self, trade_date):
        """
        从缓存加载涨停数据
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
            
        Returns:
        --------
        pandas.DataFrame or None
            涨停数据，如果缓存不存在则返回None
        """
        cache_path = self.get_limit_data_cache_path(trade_date)
        if self.cache_exists(cache_path):
            try:
                data = pd.read_csv(cache_path)
                logger.debug(f"从缓存加载涨停数据: {cache_path}")
                return data
            except Exception as e:
                logger.error(f"加载涨停数据缓存失败: {str(e)}")
                return None
        return None
    
    def save_index_data(self, trade_date, data):
        """
        保存指数数据到缓存
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
        data : pandas.DataFrame
            指数数据
            
        Returns:
        --------
        bool
            保存是否成功
        """
        try:
            cache_path = self.get_index_data_cache_path(trade_date)
            data.to_csv(cache_path, index=False)
            logger.debug(f"指数数据已缓存: {cache_path}")
            return True
        except Exception as e:
            logger.error(f"保存指数数据缓存失败: {str(e)}")
            return False
    
    def load_index_data(self, trade_date):
        """
        从缓存加载指数数据
        
        Parameters:
        -----------
        trade_date : str
            交易日期，格式为YYYYMMDD
            
        Returns:
        --------
        pandas.DataFrame or None
            指数数据，如果缓存不存在则返回None
        """
        cache_path = self.get_index_data_cache_path(trade_date)
        if self.cache_exists(cache_path):
            try:
                data = pd.read_csv(cache_path)
                logger.debug(f"从缓存加载指数数据: {cache_path}")
                return data
            except Exception as e:
                logger.error(f"加载指数数据缓存失败: {str(e)}")
                return None
        return None
    
    def preload_stock_data(self, trade_dates):
        """
        预加载多个交易日的股票数据
        优化版本：支持并行加载
        
        Parameters:
        -----------
        trade_dates : list
            交易日期列表
            
        Returns:
        --------
        dict
            预加载的股票数据字典，键为交易日期，值为股票数据DataFrame
        """
        import concurrent.futures
        import threading
        
        logger.info(f"开始预加载股票数据，共{len(trade_dates)}个交易日")
        
        # 创建线程锁，保护共享资源
        lock = threading.Lock()
        
        # 定义加载单个日期数据的函数
        def load_single_date(date):
            if date not in self.preloaded_stock_data:
                data = self.load_stock_data(date)
                if data is not None:
                    with lock:
                        self.preloaded_stock_data[date] = data
                    return True
            return False
        
        # 使用线程池并行加载数据
        loaded_count = 0
        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
            futures = {executor.submit(load_single_date, date): date for date in trade_dates}
            for future in concurrent.futures.as_completed(futures):
                if future.result():
                    loaded_count += 1
        
        logger.info(f"股票数据预加载完成，共加载{len(self.preloaded_stock_data)}个交易日的数据，本次新增{loaded_count}个")
        return self.preloaded_stock_data
    
    def preload_future_returns(self, trade_dates, n_days=5):
        """
        预加载多个交易日的未来收益率数据
        优化版本：支持并行加载
        
        Parameters:
        -----------
        trade_dates : list
            交易日期列表
        n_days : int, optional
            未来收益计算天数
            
        Returns:
        --------
        dict
            预加载的未来收益率数据字典，键为(交易日期, n_days)，值为未来收益率Series
        """
        import concurrent.futures
        import threading
        
        logger.info(f"开始预加载未来收益率数据，共{len(trade_dates)}个交易日")
        
        # 创建线程锁，保护共享资源
        lock = threading.Lock()
        
        # 定义加载单个日期数据的函数
        def load_single_date(date):
            key = (date, n_days)
            if key not in self.preloaded_future_returns:
                # 创建缓存文件名
                cache_key = f"future_returns_{date}_{n_days}"
                cache_path = os.path.join(self.cache_dir, f"{cache_key}.csv")
                
                if os.path.exists(cache_path):
                    try:
                        # 从缓存加载数据
                        sector_returns = pd.read_csv(cache_path, index_col=0, squeeze=True)
                        with lock:
                            self.preloaded_future_returns[key] = sector_returns
                        return True
                    except Exception as e:
                        logger.error(f"加载未来收益率缓存失败: {str(e)}")
            return False
        
        # 使用线程池并行加载数据
        loaded_count = 0
        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
            futures = {executor.submit(load_single_date, date): date for date in trade_dates}
            for future in concurrent.futures.as_completed(futures):
                if future.result():
                    loaded_count += 1
        
        logger.info(f"未来收益率数据预加载完成，共加载{len(self.preloaded_future_returns)}个数据，本次新增{loaded_count}个")
        return self.preloaded_future_returns
    
    def preload_index_data(self, trade_dates):
        """
        预加载多个交易日的指数数据
        优化版本：支持并行加载
        
        Parameters:
        -----------
        trade_dates : list
            交易日期列表
            
        Returns:
        --------
        dict
            预加载的指数数据字典，键为交易日期，值为指数数据DataFrame
        """
        import concurrent.futures
        import threading
        
        logger.info(f"开始预加载指数数据，共{len(trade_dates)}个交易日")
        
        # 创建线程锁，保护共享资源
        lock = threading.Lock()
        
        # 定义加载单个日期数据的函数
        def load_single_date(date):
            if date not in self.preloaded_index_data:
                data = self.load_index_data(date)
                if data is not None:
                    with lock:
                        self.preloaded_index_data[date] = data
                    return True
            return False
        
        # 使用线程池并行加载数据
        loaded_count = 0
        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
            futures = {executor.submit(load_single_date, date): date for date in trade_dates}
            for future in concurrent.futures.as_completed(futures):
                if future.result():
                    loaded_count += 1
        
        logger.info(f"指数数据预加载完成，共加载{len(self.preloaded_index_data)}个交易日的数据，本次新增{loaded_count}个")
        return self.preloaded_index_data
    
    def preload_limit_data(self, trade_dates):
        """
        预加载多个交易日的涨停数据
        优化版本：支持并行加载
        
        Parameters:
        -----------
        trade_dates : list
            交易日期列表
            
        Returns:
        --------
        dict
            预加载的涨停数据字典，键为交易日期，值为涨停数据DataFrame
        """
        import concurrent.futures
        import threading
        
        logger.info(f"开始预加载涨停数据，共{len(trade_dates)}个交易日")
        
        # 创建线程锁，保护共享资源
        lock = threading.Lock()
        
        # 定义加载单个日期数据的函数
        def load_single_date(date):
            if date not in self.preloaded_limit_data:
                data = self.load_limit_data(date)
                if data is not None:
                    with lock:
                        self.preloaded_limit_data[date] = data
                    return True
            return False
        
        # 使用线程池并行加载数据
        loaded_count = 0
        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
            futures = {executor.submit(load_single_date, date): date for date in trade_dates}
            for future in concurrent.futures.as_completed(futures):
                if future.result():
                    loaded_count += 1
        
        logger.info(f"涨停数据预加载完成，共加载{len(self.preloaded_limit_data)}个交易日的数据，本次新增{loaded_count}个")
        return self.preloaded_limit_data
    
    def get_preloaded_stock_data(self, trade_date):
        """
        获取预加载的股票数据
        
        Parameters:
        -----------
        trade_date : str
            交易日期
            
        Returns:
        --------
        pandas.DataFrame or None
            预加载的股票数据，如果未预加载则返回None
        """
        return self.preloaded_stock_data.get(trade_date)
    
    def get_preloaded_future_returns(self, trade_date, n_days=5):
        """
        获取预加载的未来收益率数据
        
        Parameters:
        -----------
        trade_date : str
            交易日期
        n_days : int, optional
            未来收益计算天数
            
        Returns:
        --------
        pandas.Series or None
            预加载的未来收益率数据，如果未预加载则返回None
        """
        return self.preloaded_future_returns.get((trade_date, n_days))
    
    def get_preloaded_index_data(self, trade_date):
        """
        获取预加载的指数数据
        
        Parameters:
        -----------
        trade_date : str
            交易日期
            
        Returns:
        --------
        pandas.DataFrame or None
            预加载的指数数据，如果未预加载则返回None
        """
        return self.preloaded_index_data.get(trade_date)
    
    def get_preloaded_limit_data(self, trade_date):
        """
        获取预加载的涨停数据
        
        Parameters:
        -----------
        trade_date : str
            交易日期
            
        Returns:
        --------
        pandas.DataFrame or None
            预加载的涨停数据，如果未预加载则返回None
        """
        return self.preloaded_limit_data.get(trade_date)
        
    def clear_cache(self, older_than_days=None):
        """
        清除缓存
        
        Parameters:
        -----------
        older_than_days : int, optional
            清除指定天数之前的缓存，如果为None则清除所有缓存
            
        Returns:
        --------
        int
            清除的文件数量
        """
        count = 0
        now = datetime.now()
        
        # 清除内存中的预加载数据
        self.preloaded_stock_data.clear()
        self.preloaded_future_returns.clear()
        self.preloaded_index_data.clear()
        self.preloaded_limit_data.clear()
        
        for directory in [self.stock_data_dir, self.trade_dates_dir, self.limit_data_dir, self.index_data_dir]:
            for filename in os.listdir(directory):
                file_path = os.path.join(directory, filename)
                
                if older_than_days is not None:
                    # 获取文件修改时间
                    file_time = datetime.fromtimestamp(os.path.getmtime(file_path))
                    # 计算文件修改时间距今的天数
                    days_old = (now - file_time).days
                    
                    if days_old >= older_than_days:
                        os.remove(file_path)
                        count += 1
                else:
                    os.remove(file_path)
                    count += 1
        
        logger.info(f"清除了{count}个缓存文件")
        return count